import gzip
import struct

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

from main_tensorflow_01_model_define import CNNModel


# 定义自定义数据集类
def load_mnist(path, kind='train'):
    labels_path = f'{path}/{kind}-labels-idx1-ubyte.gz'
    images_path = f'{path}/{kind}-images-idx3-ubyte.gz'

    with gzip.open(labels_path, 'rb') as lbpath:
        struct.unpack('>II', lbpath.read(8))
        labels = np.frombuffer(lbpath.read(), dtype=np.uint8)

    with gzip.open(images_path, 'rb') as imgpath:
        struct.unpack('>IIII', imgpath.read(16))
        images = np.frombuffer(imgpath.read(), dtype=np.uint8).reshape(len(labels), 28, 28)

    return images, labels


if __name__ == '__main__':
    # 加载数据集并进行训练
    # 定义MNIST文件对应的路径
    MNIST_FILE_PATH = 'D:/TT_WORK+/PyCharm/20250109_1_CNN/MNIST/'

    # 加载训练数据
    train_images, train_labels = load_mnist(MNIST_FILE_PATH, kind='train')
    # 加载测试数据
    test_images, test_labels = load_mnist(MNIST_FILE_PATH, kind='t10k')

    # 归一化图像数据
    train_images = train_images / 255.0
    test_images = test_images / 255.0

    # 独热编码标签
    train_labels = tf.keras.utils.to_categorical(train_labels, 10)
    test_labels = tf.keras.utils.to_categorical(test_labels, 10)

    # 创建 CNNModel 实例
    cnn_model = CNNModel()
    # 打印模型摘要
    cnn_model.summary()
    # 编译模型
    cnn_model.compile()
    # 训练模型
    history = cnn_model.train(
        train_images.reshape(-1, 28, 28, 1),
        train_labels,
        epochs=5,
        batch_size=64,
        validation_data=(test_images.reshape(-1, 28, 28, 1), test_labels)
    )

    # 评估模型
    test_loss, test_acc = cnn_model.evaluate(test_images.reshape(-1, 28, 28, 1), test_labels)
    print(f'Test accuracy: {test_acc}')

    # 保存模型
    cnn_model.save('mnist_cnn_model_local.h5')

    # 绘制训练和验证的损失曲线
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(history.history['loss'], label='Training Loss')
    plt.plot(history.history['val_loss'], label='Validation Loss')
    plt.title('Loss over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()

    # 绘制训练和验证的准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(history.history['accuracy'], label='Training Accuracy')
    plt.plot(history.history['val_accuracy'], label='Validation Accuracy')
    plt.title('Accuracy over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.show()

    # # 加载模型（可选）
    # loaded_model = CNNModel.load('mnist_cnn_model.h5')
